#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
基于微调模型的NER处理器
使用微调后的BERT模型进行命名实体识别
"""

import json
import os
import torch
from transformers import AutoTokenizer, AutoModelForTokenClassification
from typing import List, Dict, Any
import logging
import numpy as np

# 配置日志
# 确保日志目录存在
log_dir = "logs"
if not os.path.exists(log_dir):
    os.makedirs(log_dir)

# 为这个模块创建独立的logger
logger = logging.getLogger('fine_tuned_ner')
logger.setLevel(logging.INFO)

# 清除可能存在的旧处理器
for handler in logger.handlers[:]:
    logger.removeHandler(handler)

# 创建文件处理器
file_handler = logging.FileHandler(os.path.join(log_dir, "fine_tuned_ner.log"), encoding='utf-8')
file_handler.setLevel(logging.INFO)

# 创建控制台处理器
console_handler = logging.StreamHandler()
console_handler.setLevel(logging.INFO)

# 创建格式化器
formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
file_handler.setFormatter(formatter)
console_handler.setFormatter(formatter)

# 添加处理器到logger
logger.addHandler(file_handler)
logger.addHandler(console_handler)

# 防止日志传播到根logger
logger.propagate = False

class FineTunedNERProcessor:
    """
    基于微调模型的命名实体识别处理器
    """
    
    def __init__(self, model_path: str = "./fine_tuned_model"):
        """
        初始化处理器
        
        Args:
            model_path: 微调模型的路径
        """
        self.model_path = model_path
        self.tokenizer = None
        self.model = None
        self.label_list = []
        self.label_to_id = {}
        self.id_to_label = {}
        
        # 检查模型路径是否存在
        if not os.path.exists(self.model_path):
            raise FileNotFoundError(f"微调模型路径不存在: {self.model_path}")
        
        # 检查必要的模型文件是否存在
        required_files = ["config.json", "pytorch_model.bin"]
        for file in required_files:
            if not os.path.exists(os.path.join(self.model_path, file)):
                raise FileNotFoundError(f"微调模型缺少必要文件: {file}")
        
        # 尝加载微调模型
        self.load_model()
    
    def load_model(self):
        """
        加载微调模型
        """
        try:
            logger.info(f"正在加载微调模型: {self.model_path}")
            
            # 检查config.json文件
            config_path = os.path.join(self.model_path, "config.json")
            if not os.path.exists(config_path):
                raise FileNotFoundError(f"找不到配置文件: {config_path}")
            
            # 加载模型配置以检查model_type
            with open(config_path, 'r', encoding='utf-8') as f:
                config = json.load(f)
                if "model_type" not in config:
                    logger.warning("配置文件中缺少model_type键")
            
            self.tokenizer = AutoTokenizer.from_pretrained(self.model_path)
            self.model = AutoModelForTokenClassification.from_pretrained(self.model_path)
            
            # 检查CUDA是否可用并移动模型
            if torch.cuda.is_available():
                self.model = self.model.to('cuda')
                logger.info("微调模型已移动到GPU")
            
            # 加载标签映射
            label_map_path = os.path.join(self.model_path, "label_mapping.json")
            if os.path.exists(label_map_path):
                with open(label_map_path, "r", encoding="utf-8") as f:
                    label_mapping = json.load(f)
                    self.label_list = label_mapping["label_list"]
                    self.label_to_id = label_mapping["label_to_id"]
                    self.id_to_label = label_mapping["id_to_label"]
                    logger.info("成功从label_mapping.json加载标签映射")
                    logger.info(f"标签映射内容: {label_mapping}")
            else:
                logger.warning(f"标签映射文件不存在: {label_map_path}")
                # 尝试从模型配置中获取标签数量
                if hasattr(self.model.config, 'num_labels'):
                    # 创建默认标签
                    self.label_list = [f"LABEL_{i}" for i in range(self.model.config.num_labels)]
                    self.label_to_id = {label: i for i, label in enumerate(self.label_list)}
                    self.id_to_label = {i: label for label, i in self.label_to_id.items()}
                    logger.info("从模型配置创建默认标签映射")
            
            # 确保标签映射一致性
            self._ensure_label_mapping_consistency()
            
            self.model.eval()
            logger.info(f"成功加载微调模型: {self.model_path}")
            logger.info(f"标签列表: {self.label_list}")
            logger.info(f"标签到ID映射: {self.label_to_id}")
            logger.info(f"ID到标签映射: {self.id_to_label}")
            logger.info(f"标签数量: {len(self.label_list)}")
            
            # 检查CUDA状态
            if torch.cuda.is_available():
                logger.info(f"模型设备: CUDA")
                logger.info(f"CUDA设备数量: {torch.cuda.device_count()}")
                for i in range(torch.cuda.device_count()):
                    logger.info(f"  设备 {i}: {torch.cuda.get_device_name(i)}")
            else:
                logger.info("模型设备: CPU")
        except Exception as e:
            logger.error(f"加载微调模型失败: {e}")
            raise
    
    def _ensure_label_mapping_consistency(self):
        """
        确保标签映射的一致性
        """
        # 确保id_to_label映射完整
        for i, label in enumerate(self.label_list):
            if i not in self.id_to_label:
                self.id_to_label[i] = label
        
        # 确保label_to_id映射完整
        for i, label in enumerate(self.label_list):
            if label not in self.label_to_id:
                self.label_to_id[label] = i
    
    def extract_entities(self, text: str) -> List[Dict[str, Any]]:
        """
        使用微调模型提取实体
        
        Args:
            text: 输入文本
            
        Returns:
            实体列表
        """
        if not self.model or not self.tokenizer:
            raise RuntimeError("模型未加载")
        
        try:
            logger.info(f"使用微调模型识别实体，输入文本: {text}")
            
            # 对文本进行编码
            inputs = self.tokenizer(
                text,
                return_tensors="pt",
                truncation=True,
                padding=True,
                max_length=512
            )
            
            # 如果CUDA可用，将输入数据移动到GPU
            if torch.cuda.is_available():
                inputs = {k: v.to('cuda') for k, v in inputs.items()}
            
            logger.debug(f"输入编码: {inputs}")
            
            # 模型推理
            with torch.no_grad():
                outputs = self.model(**inputs)
                logits = outputs.logits
                predictions = torch.argmax(logits, dim=-1)
                # 获取预测概率
                probabilities = torch.softmax(logits, dim=-1)
            
            logger.debug(f"模型输出形状: {logits.shape}")
            logger.debug(f"预测结果形状: {predictions.shape}")
            
            # 解码预测结果
            tokens = self.tokenizer.convert_ids_to_tokens(inputs["input_ids"][0])
            
            # 安全地处理预测标签和概率
            predicted_labels = []
            predicted_probs = []
            for i, id_tensor in enumerate(predictions[0]):
                id_val = id_tensor.item()
                prob_val = torch.max(probabilities[0][i]).item()
                
                # 检查ID是否在映射范围内
                if id_val in self.id_to_label:
                    predicted_labels.append(self.id_to_label[id_val])
                else:
                    # 检查是否可以通过模型配置获取标签数量范围
                    num_labels = len(self.label_list) if self.label_list else getattr(self.model.config, 'num_labels', 0)
                    if 0 <= id_val < num_labels:
                        # 如果在范围内但不在映射中，尝试直接使用索引访问
                        try:
                            label = self.label_list[id_val]
                            predicted_labels.append(label)
                            logger.debug(f"标签ID {id_val} 不在id_to_label映射中，但可以通过索引访问，标签为: {label}")
                        except IndexError:
                            predicted_labels.append('O')
                            logger.debug(f"标签ID {id_val} 超出标签列表范围，使用默认标签'O'")
                    else:
                        # 如果ID不在范围内，使用默认标签'O'
                        predicted_labels.append('O')
                        logger.debug(f"预测标签ID {id_val} 超出范围 [0, {num_labels})，使用默认标签'O'")
                
                predicted_probs.append(prob_val)
            
            logger.info(f"Tokens: {tokens}")
            logger.info(f"Predicted labels: {predicted_labels}")
            logger.info(f"Predicted probabilities: {predicted_probs}")
            
            # 处理预测结果，提取实体
            entities = self._extract_entities_from_predictions(text, tokens, predicted_labels, predicted_probs)
            
            logger.info(f"识别到的实体: {entities}")
            return entities
            
        except Exception as e:
            logger.error(f"实体识别过程中出错: {e}")
            import traceback
            logger.error(traceback.format_exc())
            return []
    
    def _extract_entities_from_predictions(self, original_text: str, tokens: List[str], 
                                         predicted_labels: List[str], predicted_probs: List[float]) -> List[Dict[str, Any]]:
        """
        从预测结果中提取实体
        
        Args:
            original_text: 原始文本
            tokens: 分词结果
            predicted_labels: 预测标签
            predicted_probs: 预测概率
            
        Returns:
            实体列表
        """
        logger.info(f"从预测结果中提取实体，原始文本: {original_text}")
        logger.info(f"Tokens: {tokens}")
        logger.info(f"预测标签: {predicted_labels}")
        logger.info(f"预测概率: {predicted_probs}")
        
        entities = []
        current_entity = None
        current_entity_probs = []
        
        # 重构文本和标签的对应关系
        token_index = 0
        text_index = 0
        
        for i, (token, label, prob) in enumerate(zip(tokens, predicted_labels, predicted_probs)):
            logger.debug(f"处理 token: {token}, label: {label}, prob: {prob}")
            
            # 跳过特殊标记
            if token in ['[CLS]', '[SEP]', '[PAD]'] or label == -100:
                logger.debug(f"跳过特殊标记: {token}")
                continue
            
            # 处理被分割的词
            if token.startswith('##'):
                # 这是词的一部分
                token_text = token[2:]
                logger.debug(f"处理子词: {token_text}")
            else:
                token_text = token
                logger.debug(f"处理完整词: {token_text}")
            
            # 在原始文本中查找token的位置
            pos = original_text.find(token_text, text_index)
            if pos != -1:
                text_index = pos + len(token_text)
                logger.debug(f"在原始文本中找到token，位置: {pos}-{text_index}")
                
                # 设置置信度阈值，过滤低置信度的预测
                confidence_threshold = 0.3  # 进一步降低阈值以提高召回率
                if label != 'O' and prob >= confidence_threshold:
                    # 实体标签且置信度足够高
                    logger.info(f"发现高置信度实体标签: {label} (prob: {prob})")
                    if current_entity is None:
                        # 开始新的实体
                        current_entity = {
                            "nerKey": label,
                            "nerValue": token_text,
                            "start": pos,
                            "end": pos + len(token_text)
                        }
                        current_entity_probs = [prob]
                        logger.info(f"开始新实体: {current_entity}")
                    elif current_entity["nerKey"] == label:
                        # 继续当前实体
                        current_entity["nerValue"] += token_text
                        current_entity["end"] = pos + len(token_text)
                        current_entity_probs.append(prob)
                        logger.info(f"继续当前实体: {current_entity}")
                    else:
                        # 不同类型的实体，保存当前实体并开始新的
                        # 计算平均置信度
                        avg_prob = sum(current_entity_probs) / len(current_entity_probs) if current_entity_probs else 0.0
                        current_entity["confidence"] = avg_prob
                        logger.info(f"保存当前实体并开始新实体: {current_entity}")
                        entities.append(current_entity)
                        
                        current_entity = {
                            "nerKey": label,
                            "nerValue": token_text,
                            "start": pos,
                            "end": pos + len(token_text)
                        }
                        current_entity_probs = [prob]
                else:
                    # 非实体标签或置信度不够高
                    logger.debug(f"非实体标签或低置信度 (label: {label}, prob: {prob})")
                    if current_entity is not None:
                        # 计算平均置信度
                        avg_prob = sum(current_entity_probs) / len(current_entity_probs) if current_entity_probs else 0.0
                        current_entity["confidence"] = avg_prob
                        logger.info(f"保存实体: {current_entity}")
                        entities.append(current_entity)
                        current_entity = None
                        current_entity_probs = []
            else:
                logger.debug(f"在原始文本中未找到token: {token_text}")
        
        # 添加最后一个实体（如果有）
        if current_entity is not None:
            # 计算平均置信度
            avg_prob = sum(current_entity_probs) / len(current_entity_probs) if current_entity_probs else 0.0
            current_entity["confidence"] = avg_prob
            logger.info(f"保存最后一个实体: {current_entity}")
            entities.append(current_entity)
        
        # 移除位置信息并确保类型正确
        for entity in entities:
            entity.pop("start", None)
            entity.pop("end", None)
            # 确保置信度是Python原生类型
            if isinstance(entity.get("confidence"), (np.float32, np.float64)):
                entity["confidence"] = float(entity["confidence"])
        
        logger.info(f"最终实体列表: {entities}")
        return entities

# 测试代码
if __name__ == "__main__":
    print("FineTunedNERProcessor模块已创建")
    print("请先微调模型，然后使用此处理器进行实体识别")